﻿using System;
using System.Collections.Generic;
using System.Linq;
using Algoritmia.StateSpaces;
using Algoritmia.Utils;
using Lists = Algoritmia.DataStructures.Lists;
using Sets = Algoritmia.DataStructures.Sets;

namespace Algoritmia.Problems.Puzzles.NQueens
{
    public class NQueensStateSpace4 : IReversibleForwardStateSpace<NQueensStateSpace4.State,int>
    {
        protected readonly int _N;

        public class State
        {
            private readonly int _N;
            public State Prev;
            public int Row;
            public readonly int N;
            public Sets.ISet<int> AvailableRows;
            private readonly bool[] _diag;
            private readonly bool[] _gaid;

            public State(int N, State prev = null, int row = -1)
            {
                _N = N;
                Prev = prev;
                Row = row;

                if ( prev == null)
                {
                    this.N = 0;
                    AvailableRows = AlgoritmiaFactory.Sets<int>.BuildDefault(Enumerable.Range(0, N));
                    _diag = Enumerable.Repeat(true, 2*_N - 1).ToArray();
                    _gaid = Enumerable.Repeat(true, 2*_N - 1).ToArray();
                }
                else
                {
                    this.N = prev.N + 1;
                    AvailableRows = AlgoritmiaFactory.Sets<int>.BuildDefault(prev.AvailableRows);
                    _diag = prev._diag;
                    _gaid = prev._gaid;
                    var col = this.N - 1;
                    _diag[col + row] = _gaid[GaidIndex(col,row)] = false;
                    AvailableRows.Remove(row);
                }
            }

            private int GaidIndex(int col, int row)
            {
                var gaidIndex = (col - row + _N - 1);
                gaidIndex = gaidIndex < 0
                                    ? _gaid.Length + gaidIndex
                                    : gaidIndex;
                return gaidIndex;
            }

            public void RemoveRow(int row)
            {
                var col = N - 1;
                AvailableRows.Add(row);
                _gaid[col - row + _N - 1] = _diag[col + row] = true;
            }

            public bool CellIsFree(int col, int row)
            {
                return _diag[col + row] && _gaid[GaidIndex(col, row)];
            }

            public virtual IEnumerable<int> FreeRows()
            {
                var col = N;
                foreach(var row in AvailableRows)
                {
                    if (CellIsFree(col,row))
                    {
                        yield return row;
                    }
                }
            }
        }

        public NQueensStateSpace4(int N)
        {
            _N = N;
        }

        public int N
        {
            get { return _N; }
        }

        #region Implementation of IForwardStateSpace<State,int>

        public State Decide(State state, int decision)
        {
           return new State(_N, state, decision);
        }

        public IEnumerable<int> Decisions(State state)
        {
            return state.FreeRows();
        }

        public IEnumerable<State> InitialStates
        {
            get { yield return new State(_N); }
        }

        public bool IsFinal(State state)
        {
           return state.N == _N;
        }

        #endregion

        #region Implementation of IReversibleForwardStateSpace<State,int>

        public State Undo(State state, int decision)
        {
            state.RemoveRow(decision);
            return state.Prev;
        }

        #endregion
    }

    public class NQueensEnumerator4
    {
        protected readonly Func<Lists.IList<int>> _createList;

        public NQueensEnumerator4(Func<Lists.IList<int>> createList = null)
        {
            _createList = createList ?? (() => AlgoritmiaFactory.Lists<int>.Default);
        }

        protected virtual IEnumerable<Lists.IList<int>> Backtracking ( NQueensStateSpace4.State state,
            NQueensStateSpace4 space, Lists.IList<int> decisions)
        {
            if ( space.IsFinal(state))
            {
                yield return decisions;
            }
            foreach(var decision in space.Decisions(state))
            {
                decisions.Add(decision);
                state = space.Decide(state, decision);
                foreach(var result in Backtracking(state,space,decisions))
                {
                    yield return result;
                }
                decisions.RemoveAt(decisions.Count-1);
                state = space.Undo(state, decision);
            }
        }

        public virtual IEnumerable<Lists.IList<int>> Enumerate ( NQueensStateSpace4 space)
        {
            var decisions = _createList();
            foreach(var initial in space.InitialStates)
            {
                foreach(var result in Backtracking(initial,space,decisions))
                {
                    yield return result;
                }
            }
        }

        public Lists.IList<int> FirstSolution( NQueensStateSpace4 space)
        {
            return Enumerate(space).First();
        }
    }

}
